`timescale 1ns/1ps
`include "define.sv"

module systolic_system_tb;

// 时钟和复位信号
reg clk;
reg reset;

// 矩阵输入接口
reg [31:0] a_din [0:31];
reg a_wr_en [0:31];
reg [31:0] b_din [0:31];
reg b_wr_en [0:31];
reg [31:0] c_din [0:31][0:31];
reg c_wr_en [0:31][0:31];

// 控制信号
reg [1:0] matrix_type;
reg [1:0] data_type;
reg mix_precision;

// 输出接口
wire [31:0] c_out_muxed;

// 预期结果存储
reg [31:0] expected_D [0:15][0:15];

// 实例化被测模块
systolic_system uut (
    .clk(clk),
    .reset(reset),
    .a_din(a_din),
    .a_wr_en(a_wr_en),
    .b_din(b_din),
    .b_wr_en(b_wr_en),
    .c_din(c_din),
    .c_wr_en(c_wr_en),
    .matrix_type(matrix_type),
    .data_type(data_type),
    .mix_precision(mix_precision),
    .c_out_muxed(c_out_muxed)
);

// 时钟生成
initial begin
    clk = 0;
    forever #5 clk = ~clk;
end

// 测试主流程
initial begin
    reset = 1;
    matrix_type = 0;
    data_type = 2'b01;  // INT8模式
    mix_precision = 0;
    
    // 初始化使能信号
    foreach(a_wr_en[i]) a_wr_en[i] = 0;
    foreach(b_wr_en[i]) b_wr_en[i] = 0;
    foreach(c_wr_en[i,j]) c_wr_en[i][j] = 0;
    
    // 加载测试数据
    load_test_data();
    
    // 释放复位并写入数据
    #20 reset = 0;
    write_fifo_data();
    
    // 等待计算完成
    #2000;
    
    // 验证结果
    verify_results();
    $display("测试通过!");
    $finish;
end

//---------------------------------------------------------------------
// 任务：加载测试数据
//---------------------------------------------------------------------
task load_test_data();
    automatic int fd;
    automatic logic [7:0] data;
    
    // 加载A矩阵（按行存储到前16通道）
    fd = $fopen("a_int8_m16n16k16_hex.txt", "r");
    if (fd) begin
        for (int row=0; row<16; row++) begin
            for (int col=0; col<16; col++) begin
                $fscanf(fd, "%h", data);
                a_din[row][col*8 +:8] = data;  // 低8位存储
            end
            // 通道17-31填充0
            for (int col=16; col<32; col++) a_din[row][col*8 +:8] = 8'h0;
        end
        // 通道16-31填充0
        for (int ch=16; ch<32; ch++) begin
            for (int col=0; col<32; col++) a_din[ch][col*8 +:8] = 8'h0;
        end
        $fclose(fd);
    end else $error("A文件打开失败");

    // 加载B矩阵（按列循环分配至32通道）
    fd = $fopen("b_int8_m16n16k16_hex.txt", "r");
    if (fd) begin
        automatic logic [7:0] b_data [0:255];
        // 读取全部B数据
        for (int i=0; i<256; i++) $fscanf(fd, "%h", b_data[i]);
        
        // 按列循环分配至通道
        for (int col=0; col<16; col++) begin
            automatic int ch = col % 32;  // 循环分配
            for (int row=0; row<16; row++) begin
                b_din[ch][row*8 +:8] = b_data[col*16 + row];
            end
            // 通道超出部分填零
            for (int row=16; row<32; row++) b_din[ch][row*8 +:8] = 8'h0;
        end
        // 未使用的通道填零
        for (int ch=16; ch<32; ch++) begin
            for (int row=0; row<32; row++) b_din[ch][row*8 +:8] = 8'h0;
        end
        $fclose(fd);
    end else $error("B文件打开失败");

    // 加载C矩阵（映射到32x32的左上角16x16）
    fd = $fopen("c_int8_m16n16k16_hex.txt", "r");
    if (fd) begin
        for (int i=0; i<16; i++) begin
            for (int j=0; j<16; j++) begin
                $fscanf(fd, "%h", c_din[i][j][7:0]);
                c_din[i][j][31:8] = 24'h0;  // 高24位填零
            end
            // 行16-31填零
            for (int j=16; j<32; j++) c_din[i][j] = 32'h0;
        end
        // 行16-31填零
        for (int i=16; i<32; i++) begin
            for (int j=0; j<32; j++) c_din[i][j] = 32'h0;
        end
        $fclose(fd);
    end else $error("C文件打开失败");
endtask

//---------------------------------------------------------------------
// 任务：写入FIFO数据
//---------------------------------------------------------------------
task write_fifo_data();
    // 写入A矩阵（按行顺序写入前16通道）
    for (int cycle=0; cycle<16; cycle++) begin
        @(posedge clk);
        foreach(a_wr_en[i]) begin
            a_wr_en[i] = (i < 16) ? 1'b1 : 1'b0;  // 前16通道使能
            a_din[i] = {24'h0, a_din[i][cycle*8 +:8]};  // 每个周期写入一行
        end
        #1;
    end
    @(posedge clk);
    foreach(a_wr_en[i]) a_wr_en[i] = 0;

    // 写入B矩阵（按列间隔写入）
    for (int cycle=0; cycle<16; cycle++) begin
        @(posedge clk);
        foreach(b_wr_en[i]) begin
            b_wr_en[i] = (i < 16) ? 1'b1 : 1'b0;  // 前16通道使能
            b_din[i] = {24'h0, b_din[i][cycle*8 +:8]};  // 每个周期写入一列
        end
        #1;
    end
    @(posedge clk);
    foreach(b_wr_en[i]) b_wr_en[i] = 0;

    // 写入C矩阵（一次性写入前16x16区域）
    @(posedge clk);
    foreach(c_wr_en[i,j]) begin
        c_wr_en[i][j] = (i < 16 && j < 16) ? 1'b1 : 1'b0;
    end
    #1;
    @(posedge clk);
    foreach(c_wr_en[i,j]) c_wr_en[i][j] = 0;
endtask

//---------------------------------------------------------------------
// 任务：验证结果
//---------------------------------------------------------------------
task verify_results();
    reg [31:0] captured_c_out [0:15][0:15];
    automatic integer cycle_count = 0;

    // 捕获256个周期的复用输出
    fork
        begin
            while (cycle_count < 256) begin
                @(posedge clk);
                captured_c_out[uut.counter[7:4]][uut.counter[3:0]] = c_out_muxed;
                cycle_count++;
            end
        end
        // 超时保护
        #5000 $error("超时：未捕获到全部输出");
    join

    // 与预期结果比较
    for (int i=0; i<16; i++) begin
        for (int j=0; j<16; j++) begin
            if (captured_c_out[i][j] !== expected_D[i][j]) begin
                $error("PE[%0d][%0d] 错误：预期值=%h，实际值=%h", i, j, expected_D[i][j], captured_c_out[i][j]);
            end
        end
    end
endtask

//---------------------------------------------------------------------
// 加载预期结果（D矩阵）
//---------------------------------------------------------------------
initial begin
    for (int i=0; i<16; i++) begin
        automatic string d_file = $sformatf("D_data_row%0d.txt", i);
        $readmemh(d_file, expected_D[i]);
    end
end

endmodule